package com.monitor.util;

import com.jcraft.jsch.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Properties;

/**
 * SSH客户端工具类
 * 提供SSH连接和命令执行功能，支持密码和密钥认证
 */
public class SSHClient {
    private static final Logger logger = LoggerFactory.getLogger(SSHClient.class);
    private JSch jsch;           // JSch库实例
    private Session session;     // SSH会话
    private int connectionTimeout; // 连接超时时间(毫秒)
    private int commandTimeout;    // 命令执行超时时间(毫秒)

    /**
     * 构造SSH客户端
     * @param connectionTimeout 连接超时时间(毫秒)
     * @param commandTimeout 命令执行超时时间(毫秒)
     */
    public SSHClient(int connectionTimeout, int commandTimeout) {
        this.jsch = new JSch();
        this.connectionTimeout = connectionTimeout;
        this.commandTimeout = commandTimeout;
    }

    /**
     * 使用密码方式连接SSH服务器
     * @param host 主机地址
     * @param port SSH端口
     * @param username 用户名
     * @param password 密码
     * @return 连接是否成功
     */
    public boolean connect(String host, int port, String username, String password) {
        try {
            // 创建SSH会话
            session = jsch.getSession(username, host, port);
            session.setPassword(password);
            
            // 配置会话参数
            Properties config = new Properties();
            // 禁用严格主机密钥检查，避免首次连接时询问
            config.put("StrictHostKeyChecking", "no");
            session.setConfig(config);
            session.setTimeout(connectionTimeout);
            
            logger.info("正在连接到服务器: {}@{}:{}", username, host, port);
            // 建立连接
            session.connect(connectionTimeout);
            logger.info("成功连接到服务器: {}@{}:{}", username, host, port);
            return true;
        } catch (JSchException e) {
            logger.error("连接服务器失败: {}@{}:{}, 错误: {}", username, host, port, e.getMessage());
            return false;
        }
    }

    /**
     * 使用密钥方式连接SSH服务器
     * @param host 主机地址
     * @param port SSH端口
     * @param username 用户名
     * @param privateKeyPath 私钥文件路径
     * @return 连接是否成功
     */
    public boolean connectWithKey(String host, int port, String username, String privateKeyPath) {
        try {
            // 加载私钥
            jsch.addIdentity(privateKeyPath);
            // 创建SSH会话
            session = jsch.getSession(username, host, port);
            
            // 配置会话参数
            Properties config = new Properties();
            // 禁用严格主机密钥检查，避免首次连接时询问
            config.put("StrictHostKeyChecking", "no");
            session.setConfig(config);
            session.setTimeout(connectionTimeout);
            
            logger.info("正在使用密钥连接到服务器: {}@{}:{}", username, host, port);
            // 建立连接
            session.connect(connectionTimeout);
            logger.info("成功使用密钥连接到服务器: {}@{}:{}", username, host, port);
            return true;
        } catch (JSchException e) {
            logger.error("使用密钥连接服务器失败: {}@{}:{}, 错误: {}", username, host, port, e.getMessage());
            return false;
        }
    }

    /**
     * 在SSH服务器上执行命令
     * @param command 要执行的命令
     * @return 命令执行结果
     */
    public String executeCommand(String command) {
        Channel channel = null;
        try {
            logger.debug("执行命令: {}", command);
            // 创建exec类型的通道用于执行命令
            channel = session.openChannel("exec");
            ((ChannelExec) channel).setCommand(command);
            
            // 设置标准输出和错误输出的处理
            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
            channel.setOutputStream(outputStream);
            ((ChannelExec) channel).setErrStream(outputStream);
            
            // 获取输入流，用于读取命令执行的结果
            InputStream in = channel.getInputStream();
            // 连接通道并开始执行命令
            channel.connect(commandTimeout);
            
            byte[] tmp = new byte[1024];
            StringBuilder result = new StringBuilder();
            
            // 等待命令执行完成并读取结果
            while (true) {
                // 读取可用的输出数据
                while (in.available() > 0) {
                    int i = in.read(tmp, 0, 1024);
                    if (i < 0) break;
                    result.append(new String(tmp, 0, i));
                }
                // 如果通道已关闭且没有更多数据可读，则退出循环
                if (channel.isClosed()) {
                    if (in.available() > 0) continue;
                    break;
                }
                try {
                    // 短暂等待，避免CPU占用过高
                    Thread.sleep(100);
                } catch (Exception ee) {
                    // 忽略中断异常
                }
            }
            
            // 追加任何额外的输出
            result.append(outputStream.toString());
            int exitStatus = channel.getExitStatus();
            logger.debug("命令执行完成，退出状态码: {}", exitStatus);
            return result.toString();
        } catch (JSchException | IOException e) {
            logger.error("执行命令失败: {}, 错误: {}", command, e.getMessage());
            return "ERROR: " + e.getMessage();
        } finally {
            // 确保通道被关闭
            if (channel != null) {
                channel.disconnect();
            }
        }
    }

    /**
     * 断开SSH连接
     */
    public void disconnect() {
        if (session != null && session.isConnected()) {
            logger.debug("断开SSH连接");
            session.disconnect();
        }
    }
} 